import akshare as ak
import pandas as pd
from pyecharts import options as opts
from pyecharts.charts import Kline, Line
import numpy as np

def get_shanghai_index_data(start_date, end_date):
    """
    获取指定日期范围内的上证指数日线数据

    :param start_date: 起始日期，格式为 'YYYY-MM-DD'
    :param end_date: 结束日期，格式为 'YYYY-MM-DD'
    :return: 筛选后的上证指数日线数据
    """
    # 获取上证指数日线数据
    stock_zh_index_daily_df = ak.stock_zh_index_daily(symbol="sh000001")
    # 将 date 列转换为日期时间类型
    stock_zh_index_daily_df['date'] = pd.to_datetime(stock_zh_index_daily_df['date'])
    # 根据起始日期和结束日期筛选数据
    filtered_data = stock_zh_index_daily_df[
        (stock_zh_index_daily_df['date'] >= start_date) &
        (stock_zh_index_daily_df['date'] <= end_date)
    ]
    return filtered_data

def draw_kline_and_ma_chart(data):
    """
    绘制上证指数 K 线图及 MA5、MA20 均线图

    :param data: 包含上证指数数据的 DataFrame
    :return: 合并后的 K 线图和均线图对象
    """
    # 准备 K 线图所需的数据，并保留两位小数
    kline_data = [[round(val, 2) for val in row] for row in data[['open', 'close', 'low', 'high']].values.tolist()]
    # 从 date 列获取日期数据并转换为字符串列表
    x_data = data['date'].dt.strftime('%Y-%m-%d').tolist()
    close_prices = data['close'].values

    # 创建 K 线图对象，设置宽度和高度扩大两倍
    kline = (
        Kline(init_opts=opts.InitOpts(width="1200px", height="600px"))
        .add_xaxis(x_data)
        .add_yaxis(
            "上证指数",
            kline_data,
            itemstyle_opts=opts.ItemStyleOpts(
                color="#ec0000",  # 阳线颜色
                color0="#00da3c",  # 阴线颜色
                border_color="#ec0000",
                border_color0="#00da3c"
            )
        )
        .set_global_opts(
            title_opts=opts.TitleOpts(
                title="上证指数K线图",
                pos_left="10%"  # 将标题设置为距离画布左侧 5% 的位置
            ),
            xaxis_opts=opts.AxisOpts(
                is_scale=True,
                # 不显示 x 轴轴线
                axisline_opts=opts.AxisLineOpts(is_show=False),
                # 不显示 x 轴刻度
                axistick_opts=opts.AxisTickOpts(is_show=False),
                # 不显示 x 轴标签
                axislabel_opts=opts.LabelOpts(is_show=False),
                # 不显示 x 轴分割线
                splitline_opts=opts.SplitLineOpts(is_show=False),
            ),
            yaxis_opts=opts.AxisOpts(
                is_scale=True,
                # 不显示 y 轴轴线
                axisline_opts=opts.AxisLineOpts(is_show=False),
                # 不显示 y 轴刻度
                axistick_opts=opts.AxisTickOpts(is_show=False),
                # 不显示 y 轴标签
                axislabel_opts=opts.LabelOpts(is_show=False),
                # 不显示 y 轴分割线
                splitline_opts=opts.SplitLineOpts(is_show=False),
                splitarea_opts=opts.SplitAreaOpts(
                    is_show=True,
                    areastyle_opts={"color": ["#0f3c68", "#0f3c68"]}
                )
            ),
            datazoom_opts=[opts.DataZoomOpts(pos_bottom="5%")],  # 数据缩放工具
            tooltip_opts=opts.TooltipOpts(
                trigger="axis",
                axis_pointer_type="cross"  # 十字光标
            )
        )
    )

    # 添加均线（MA5和MA20），并保留两位小数
    periods = len(close_prices)
    ma5 = [round(np.mean(close_prices[max(0, i - 5):i + 1]), 2) for i in range(periods)]
    ma20 = [round(np.mean(close_prices[max(0, i - 20):i + 1]), 2) for i in range(periods)]
    line = (
        Line()
        .add_xaxis(x_data)
        .add_yaxis(
            "MA5",
            ma5,
            is_smooth=True,
            linestyle_opts=opts.LineStyleOpts(width=2),
            # 不显示 MA5 的数值标签
            label_opts=opts.LabelOpts(is_show=False)
        )
        .add_yaxis(
            "MA20",
            ma20,
            is_smooth=True,
            linestyle_opts=opts.LineStyleOpts(width=2),
            # 不显示 MA20 的数值标签
            label_opts=opts.LabelOpts(is_show=False)
        )
    )

    # 合并图表
    combined_chart = kline.overlap(line)
    return combined_chart

if __name__ == "__main__":
    # 自定义起始日期和结束日期
    start_date = "2025-01-01"
    end_date = "2025-06-01"
    # 获取数据
    shanghai_data = get_shanghai_index_data(start_date, end_date)
    # 绘制 K 线图和均线图
    chart = draw_kline_and_ma_chart(shanghai_data)
    # 渲染为 HTML 文件
    chart.render("ssec_kline.html")